package com.xagent.engine.policy;

import com.xagent.core.Properties;
import com.xagent.engine.context.AgentContext;
import com.xagent.engine.hook.HookInfo;
import com.xagent.engine.hook.MethodInfo;
import com.xagent.engine.hook.impl.HookInfoImpl;
import com.xagent.engine.hook.impl.MethodInfoImpl;
import com.xagent.engine.utils.HookUtil;
import org.slf4j.Logger;
import org.xml.sax.Attributes;
import org.xml.sax.SAXException;
import org.xml.sax.helpers.DefaultHandler;

import javax.xml.parsers.ParserConfigurationException;
import javax.xml.parsers.SAXParser;
import javax.xml.parsers.SAXParserFactory;
import java.io.*;
import java.util.HashMap;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class PolicyBuilder {
    public static void parsePolicyByInputStream(InputStream inputStream, HashMap<String, HashMap<String, HookInfo>> hooks, String type) {
        SAXParserFactory factory = SAXParserFactory.newInstance();
        SAXParser parser = null;
        try {
            parser = factory.newSAXParser();
            parser.parse(inputStream, new PolicyHandler(hooks, type));
        } catch (ParserConfigurationException | SAXException | IOException e) {
            //

        }
    }

}

class PolicyHandler extends DefaultHandler {
    private boolean inMethod = false;
    private HashMap<String, HashMap<String, HookInfo>> hooks;
    private String type = "";

    public PolicyHandler(HashMap<String, HashMap<String, HookInfo>> hooks, String type) {
        this.hooks = hooks;
        this.type = type;
    }

    @Override
    public void startElement(String uri, String localName, String qName, Attributes attributes) throws SAXException {
        if (qName.equalsIgnoreCase("method")) {
            Logger agentLogger = AgentContext.getContext().getAgentLogger();
            Logger securityLogger = AgentContext.getContext().getSecurityLogger();
            inMethod = true;
            String id = attributes.getValue("id");
            String signature = attributes.getValue("signature");
            String source = attributes.getValue("source");
            String target = attributes.getValue("target");
            String index = attributes.getValue("index");
            String className = "";
            String methodName = "";
            String methodDescriptor = "";
            Matcher matcher = Pattern.compile("([a-zA-Z\\.]*)(\\([a-zA-Z\\.;/\\[\\],]*\\))").matcher(signature); // 没匹配返回值
            if (matcher.find()) {
                className = matcher.group(1).substring(0, matcher.group(1).lastIndexOf("."));
                methodName = matcher.group(1).substring(matcher.group(1).lastIndexOf(".") + 1, matcher.group(1).length());
                methodDescriptor = HookUtil.standardizationSignature(matcher.group(2));
            }

            if (!checkIfHooked(hooks.get(type), className, methodName, methodDescriptor.toString())) {
                MethodInfo methodInfo = new MethodInfoImpl(id, className, methodName, methodDescriptor.toString(), source, target);
                methodInfo.setIndex(Integer.parseInt(index));
                hooks.get(type).get(className).addMethodInfo(methodInfo);

                if (AgentContext.getContext().isDebug()) {
                    agentLogger.debug("[DEBUG] " + "Recording " + className + ":" + methodName + "," + " VulnTypeId:" + id + " descriptor:" + methodDescriptor);
                }

            }
        }


    }

    @Override
    public void endElement(String uri, String localName, String qName) throws SAXException {
        if (qName.equalsIgnoreCase("method")) {
            inMethod = false;
        }
    }
//    @Override
//    public void characters(char[] ch, int start, int length) throws SAXException {
//        super.characters(ch, start, length);
//    }

    /**
     * @param hooks            hook信息Map
     * @param className        类名
     * @param methodName       方法名
     * @param methodDescriptor 方法描述符
     * @return 当前方法是否需要hook，true代表需要hook，false代表不hook.
     */
    public boolean checkIfHooked(HashMap<String, HookInfo> hooks, String className, String methodName, String methodDescriptor) {
        boolean flag = false;
        if (!hooks.containsKey(className)) {
            hooks.put(className, new HookInfoImpl(className));
        }
        HookInfoImpl hookInfoImpl = (HookInfoImpl) hooks.get(className);
        if (hookInfoImpl.getMethodInfoHashMap().containsKey(methodName + methodDescriptor)) {
            flag = true;
        }

        return flag;
    }


}

